#!/usr/bin/python
'''
Reads in a fasta file and prints out simplified names and a codex to translate.
'''
import sys

# parameters.
in_file = sys.argv[1]
agp_file = sys.argv[2]
out_fasta = sys.argv[3]
out_agp = sys.argv[4]
codec_file = sys.argv[5]
ncbi_file = sys.argv[6]

# read in the ncbi file.
skip = {}
fin = open(ncbi_file, "rb")
for line in fin:
	tmp = line.strip()
	if len(tmp) > 5:
		skip[tmp] = True
fin.close()
'''
# open outputs.
fout = open(out_fasta, "wb")
cout = open(codec_file, "wb")

# parse fasta.
idx = 0
j = 0
trans = {}
first = True
fin = open(in_file, "rb")
for line in fin:
	# check if we start new entry.
	if line[0] == ">":
		# tally.
		if first != True:
			# See if we keep this read.
			if head not in skip:
				# make new name.
				oname = head
				nname = "contig%i" % idx
				idx += 1
		
				# mark.
				trans[oname] = nname
				
				# write.
				fout.write(">%s\n" % nname)
				fout.write("%s\n" % seq)
				
				# write codec.
				cout.write("%s\t%s\n" % (oname, nname))
			else:
				print j
				j += 1
				skip[head] = False
		else:
			first = False
			
		# clear.
		head = line.strip().replace(">","")
		seq = ""
		continue
	

	# get seq.
	seq += line.strip()

# get last one.
if head not in skip:
	# make new name.
	oname = head
	nname = "contig%i" % idx
	idx += 1

	# mark.
	trans[oname] = nname
	
	# write.
	fout.write(">%s\n" % nname)
	fout.write("%s\n" % seq)
	
	# write codec.
	cout.write("%s\t%s\n" % (oname, nname))
fin.close()
fout.close()
cout.close()
'''

# read in agp.
scafs = {}
fin = open(agp_file, "rb")
i = 0
for line in fin:
	# split line.
	tmp = line.strip().split("\t")
	
	# check if scaffold exists.
	if tmp[0] not in scafs:
		scafs[tmp[0]] = []
	
	# skip if in skips and in fasta.
	#if tmp[5] not in skip and tmp[5] in trans:
		# add good entries.
	scafs[tmp[0]].append(tmp)
		
	i += 1
#	if i > 20000: break
fin.close()

# remove stuff from agp.
fout = open(out_agp, "wb")
for key in scafs:
	# loop over scaf to find hit.
	scaf = scafs[key]
	
	# remove fragments at start.
	tmp = []
	noskip = False
	for i in range(len(scaf)):
		# stop when contig found.
		if scaf[i][4] == "W":
			noskip = True
			
		# just add everything when noskip.
		if noskip == True:
			tmp.append(scaf[i])
			continue
			
	
	# remove fragments from end.
	tmp1 = []
	noskip = False
	for i in range(len(tmp)-1, -1, -1):
		# stop when contig found.
		if tmp[i][4] == "W":
			noskip = True
			
		# just add everything when noskip.
		if noskip == True:
			tmp1.append(tmp[i])
			continue
	
	# replace original with trimmed.
	debug = scaf
	scaf = tmp1
	scaf.reverse()
	
	# manage numbers.
	idx = 1
	for i in range(len(scaf)):
		# handle contig.
		if scaf[i][4] == "W":
			# set start.
			scaf[i][1] = str(idx)
			
			# set stop.
			sz = int(scaf[i][7])
			stop = idx + sz
			scaf[i][2] = str(stop)
			
			# update size.
			scaf[i][7] = str(sz + 1)
			
			# replace contig name.
			if scaf[i][5] not in trans:
				print "missing name"
				print scaf[i]
				sys.exit()
				
			scaf[i][5] = trans[scaf[i][5]]
			
		# handle gap.
		else:
			# set start.
			scaf[i][1] = str(idx)
	
			# set stop.
			sz = int(scaf[i][5])
			stop = idx + sz
			scaf[i][2] = str(stop)
			
		# update index.
		idx = stop + 1

	if len(debug) == 0:
		continue
	
	# write scaf.
	for s in scaf:
		fout.write('\t'.join(s) + "\n")
		
fout.close()
